import cv2
from tqdm import tqdm
import numpy as np
from PIL import Image
from transparent_background import Remover

def process_video_with_inspyrenet(args):
    """
    使用 InSPyReNet (通过 transparent-background 库) 处理视频，
    移除背景并生成一个前景位于纯黑背景上的视频。
    """
    print("正在准备 InSPyReNet 模型...")
    # 初始化 Remover，可以指定不同的模式，例如 'base' 或 'fast'
    # device='cuda:0' or 'cpu'
    remover = Remover(mode='base', device='cuda:0')

    # 准备视频读取
    cap = cv2.VideoCapture(args.input_video)
    if not cap.isOpened():
        print(f"错误: 无法打开视频文件 {args.input_video}")
        return

    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)
    if not fps or fps <= 0:
        print("警告: 无法读取有效帧率, 将使用默认值 30 FPS。")
        fps = 30

    print(f"输入视频: {width}x{height}, {fps:.2f} FPS, 共 {total_frames} 帧。")

    # 准备视频写入
    writer = cv2.VideoWriter(
        args.output_file,
        cv2.VideoWriter_fourcc(*'mp4v'),
        fps,
        (width, height)
    )
    if not writer.isOpened():
        print(f"错误: 无法创建视频写入器，请检查路径 '{args.output_file}' 是否有效。")
        cap.release()
        return

    # 逐帧处理
    print("开始使用 InSPyReNet 处理视频帧...")
    for _ in tqdm(range(total_frames), desc="InSPyReNet 处理中"):
        ret, frame = cap.read()
        if not ret:
            break
        
        # OpenCV 读取的帧是 BGR 格式，需要转换为 PIL Image 的 RGB 格式
        img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        
        # 使用 transparent-background 移除背景, 返回的是 RGBA 格式的 PIL Image
        output_image = remover.process(img, type='rgba')
        
        # 将返回的 PIL Image 转换为 numpy array 以便用 OpenCV 处理
        output_data = np.array(output_image)
        
        # 将返回的 RGBA numpy array 转换回 OpenCV 的 BGR 格式以便写入
        # 首先，分离 Alpha 通道
        alpha = output_data[:, :, 3]
        # 获取 RGB 通道
        rgb_frame = output_data[:, :, :3]
        # 将 RGB 转换为 BGR
        bgr_frame = cv2.cvtColor(rgb_frame, cv2.COLOR_RGB2BGR)
        
        # 将透明区域设置为黑色
        bgr_frame[alpha == 0] = [0, 0, 0]

        writer.write(bgr_frame)

    # 清理和收尾
    cap.release()
    writer.release()

    print(f"\n处理完成！视频已保存至: {args.output_file}")


if __name__ == "__main__":
    # 直接设置参数
    class Args:
        def __init__(self):
            self.input_video = "C:/Users/Missi/Videos/素材/nina.mp4"  # 修改为你的输入视频路径
            self.output_file = "C:/Users/Missi/Videos/素材/nina_out.mp4"  # 修改为你的输出视频路径
    
    args = Args()
    process_video_with_inspyrenet(args) 